# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
将 Cifar10 下载数据转换成图片保存至文件夹
下载路径 http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz

# 当前项目: Cifar10-Classification
# 创建时间: 2019/6/22 12:00
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import os
import cv2
import numpy as np
import pickle


# Pickle 文件读取
def load_file(filename):
    with open(filename, 'rb') as fo:
        data = pickle.load(fo, encoding='latin1')
    return data


# pickle 文件转换
def main(train_data_path, save_data_path):
    """
    :param train_data_path:   cifar 10 下载文件保存路径
    :param save_data_path:    转换图片保存路径
    :return:
    """
    # 图片标签类别
    data_label = load_file(os.path.join(train_data_path, 'batches.meta'))
    # 训练图片保存路径
    train_data_folder = os.path.join(save_data_path, 'train_data')
    os.makedirs(train_data_folder, exist_ok=True)

    # 测试图片保存路径
    test_data_folder = os.path.join(save_data_path, 'test_data')
    os.makedirs(test_data_folder, exist_ok=True)

    # ##########################################################################
    # 训练数据图片列表保存
    train_f_csv = open(os.path.join(save_data_path, 'train_data.csv'), 'w', encoding='utf-8')
    train_f_csv.write('{},{}\n'.format('TrainImage', 'GroundTruth'))

    num = 0
    for data_batch_i in os.listdir(train_data_path):

        if data_batch_i.startswith('data_batch_'):
            print('\n', data_batch_i)

            data = load_file(os.path.join(train_data_path, data_batch_i))

            for i, label_i in enumerate(data['labels']):
                print('\r{}   {}  '.format(i, label_i), end='     ')
                image = data['data'][i].reshape((3, 32, 32))
                # 维度转换
                image = np.transpose(image, [1, 2, 0])

                num += 1
                img_path = os.path.join(train_data_folder,
                                        '{}_{:05d}.bmp'.format(data_label['label_names'][label_i], num))
                cv2.imencode(".bmp", image)[1].tofile(img_path)

                train_f_csv.write('{},{}\n'.format(os.path.abspath(img_path), data_label['label_names'][label_i]))

    # ##########################################################################
    # 测试数据图片列表保存
    test_f_csv = open(os.path.join(save_data_path, 'test_data.csv'), 'w', encoding='utf-8')
    test_f_csv.write('{},{}\n'.format('TestImage', 'GroundTruth'))

    num = 0
    data = load_file(os.path.join(train_data_path, 'test_batch'))
    print('\ntest_batch')
    for i, label_i in enumerate(data['labels']):
        print('\r{}   {}  '.format(i, label_i), end='     ')
        image = data['data'][i].reshape((3, 32, 32))
        # 维度转换
        image = np.transpose(image, [1, 2, 0])

        num += 1
        img_path = os.path.join(test_data_folder,
                                'Test_{}_{:05d}.bmp'.format(data_label['label_names'][label_i], num))
        cv2.imencode(".bmp", image)[1].tofile(img_path)

        test_f_csv.write('{},{}\n'.format(os.path.abspath(img_path), data_label['label_names'][label_i]))


if __name__ == '__main__':
    # 下载数据路径
    train_path = r'D:\Desktop\Cifar10-Classification\Data\cifar-10-batches-py'
    # 保存图片路径
    save_path = os.path.join(os.path.dirname(__file__), 'Data', 'Image')

    main(train_path, save_path)
